import tkinter as tk
from tkinter import ttk, messagebox, scrolledtext, filedialog
import paramiko
import threading
import time
import socket
import os
import re
from datetime import datetime

class XShellLikeSSHClient:
    def __init__(self, root):
        self.root = root
        self.root.title("SecureBridge")
        self.root.geometry("900x700")
        self.root.configure(bg='#2b2b2b')
        
        # 设置主题颜色 (模仿Xshell的深色主题)
        self.style = ttk.Style()
        self.style.theme_use('clam')
        self.configure_styles()
        
        # SSH客户端
        self.ssh_client = None
        self.shell = None
        self.connected = False
        self.session_name = "会话1"
        
        # 创建界面
        self.create_menu()
        self.create_widgets()
        
        # 会话管理
        self.sessions = {}
        self.current_session = None
        
    def configure_styles(self):
        """配置样式模仿Xshell的深色主题"""
        # 配置颜色
        bg_color = '#2b2b2b'
        fg_color = '#ffffff'
        accent_color = '#007acc'
        panel_bg = '#3c3c3c'
        entry_bg = '#1e1e1e'
        
        self.style.configure('TFrame', background=bg_color)
        self.style.configure('TLabel', background=bg_color, foreground=fg_color, font=('Segoe UI', 9))
        self.style.configure('TButton', background=panel_bg, foreground=fg_color, font=('Segoe UI', 9))
        self.style.configure('TEntry', fieldbackground=entry_bg, foreground=fg_color, insertcolor=fg_color)
        self.style.configure('TCombobox', fieldbackground=entry_bg, foreground=fg_color, background=panel_bg)
        self.style.configure('TLabelframe', background=bg_color, foreground=accent_color)
        self.style.configure('TLabelframe.Label', background=bg_color, foreground=accent_color)
        
        # 自定义样式
        self.style.configure('Accent.TButton', background=accent_color, foreground='white')
        self.style.configure('Success.TLabel', background=bg_color, foreground='#4caf50')
        self.style.configure('Error.TLabel', background=bg_color, foreground='#f44336')
        self.style.configure('Terminal.TFrame', background='black')
        
    def create_menu(self):
        """创建菜单栏"""
        menubar = tk.Menu(self.root, bg='#2b2b2b', fg='white', activebackground='#3c3c3c')
        
        # 文件菜单
        file_menu = tk.Menu(menubar, tearoff=0, bg='#2b2b2b', fg='white')
        file_menu.add_command(label="新建会话", command=self.new_session, accelerator="Ctrl+N")
        file_menu.add_command(label="打开会话", command=self.open_session, accelerator="Ctrl+O")
        file_menu.add_command(label="保存会话", command=self.save_session, accelerator="Ctrl+S")
        file_menu.add_separator()
        file_menu.add_command(label="退出", command=self.root.quit)
        menubar.add_cascade(label="文件", menu=file_menu)
        
        # 编辑菜单
        edit_menu = tk.Menu(menubar, tearoff=0, bg='#2b2b2b', fg='white')
        edit_menu.add_command(label="复制", command=self.copy_text, accelerator="Ctrl+C")
        edit_menu.add_command(label="粘贴", command=self.paste_text, accelerator="Ctrl+V")
        edit_menu.add_separator()
        edit_menu.add_command(label="全选", command=self.select_all, accelerator="Ctrl+A")
        menubar.add_cascade(label="编辑", menu=edit_menu)
        
        # 查看菜单
        view_menu = tk.Menu(menubar, tearoff=0, bg='#2b2b2b', fg='white')
        view_menu.add_command(label="切换主题", command=self.toggle_theme)
        view_menu.add_command(label="字体设置", command=self.font_settings)
        view_menu.add_separator()
        view_menu.add_command(label="清除终端", command=self.clear_screen)
        menubar.add_cascade(label="查看", menu=view_menu)
        
        # 工具菜单
        tools_menu = tk.Menu(menubar, tearoff=0, bg='#2b2b2b', fg='white')
        tools_menu.add_command(label="会话管理器", command=self.session_manager)
        tools_menu.add_command(label="选项", command=self.options)
        menubar.add_cascade(label="工具", menu=tools_menu)
        
        self.root.config(menu=menubar)
        
        # 绑定快捷键
        self.root.bind('<Control-n>', lambda e: self.new_session())
        self.root.bind('<Control-o>', lambda e: self.open_session())
        self.root.bind('<Control-s>', lambda e: self.save_session())
        
    def create_widgets(self):
        """创建主界面组件"""
        # 主容器
        main_container = ttk.Frame(self.root)
        main_container.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 顶部连接栏
        self.create_connection_bar(main_container)
        
        # 终端区域
        self.create_terminal_area(main_container)
        
        # 状态栏
        self.create_status_bar(main_container)
        
    def create_connection_bar(self, parent):
        """创建连接控制栏"""
        conn_frame = ttk.LabelFrame(parent, text="会话连接", padding="10")
        conn_frame.pack(fill=tk.X, pady=(0, 5))
        
        # 第一行
        row1 = ttk.Frame(conn_frame)
        row1.pack(fill=tk.X, pady=2)
        
        ttk.Label(row1, text="主机:").pack(side=tk.LEFT, padx=(0, 5))
        self.host_entry = ttk.Entry(row1, width=20)
        self.host_entry.pack(side=tk.LEFT, padx=(0, 15))
        # 清空默认值，让用户自己填写
        self.host_entry.insert(0, "")
        
        ttk.Label(row1, text="端口:").pack(side=tk.LEFT, padx=(0, 5))
        self.port_entry = ttk.Entry(row1, width=8)
        self.port_entry.pack(side=tk.LEFT, padx=(0, 15))
        self.port_entry.insert(0, "22")
        
        # 第二行
        row2 = ttk.Frame(conn_frame)
        row2.pack(fill=tk.X, pady=2)
        
        ttk.Label(row2, text="用户名:").pack(side=tk.LEFT, padx=(0, 5))
        self.username_entry = ttk.Entry(row2, width=20)
        self.username_entry.pack(side=tk.LEFT, padx=(0, 15))
        # 清空默认值
        self.username_entry.insert(0, "")
        
        ttk.Label(row2, text="认证方法:").pack(side=tk.LEFT, padx=(0, 5))
        self.auth_method = ttk.Combobox(row2, width=15, values=["Password", "Public Key", "Keyboard-Interactive"])
        self.auth_method.pack(side=tk.LEFT, padx=(0, 15))
        self.auth_method.set("Password")
        
        ttk.Label(row2, text="密码:").pack(side=tk.LEFT, padx=(0, 5))
        self.password_entry = ttk.Entry(row2, show="•", width=20)
        self.password_entry.pack(side=tk.LEFT, padx=(0, 15))
        
        # 第三行 - 按钮行
        row3 = ttk.Frame(conn_frame)
        row3.pack(fill=tk.X, pady=(10, 0))
        
        self.connect_btn = ttk.Button(row3, text="连接", command=self.toggle_connection, style='Accent.TButton')
        self.connect_btn.pack(side=tk.LEFT, padx=(0, 10))
        
        ttk.Button(row3, text="新建会话", command=self.new_session).pack(side=tk.LEFT, padx=(0, 10))
        ttk.Button(row3, text="保存会话", command=self.save_session).pack(side=tk.LEFT, padx=(0, 10))
        ttk.Button(row3, text="断开连接", command=self.disconnect).pack(side=tk.LEFT, padx=(0, 10))
        
        # 连接状态
        self.status_label = ttk.Label(row3, text="未连接", style='Error.TLabel')
        self.status_label.pack(side=tk.RIGHT)
        
    def create_terminal_area(self, parent):
        """创建终端显示区域"""
        terminal_frame = ttk.LabelFrame(parent, text="终端", padding="5")
        terminal_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 5))
        
        # 终端工具栏
        toolbar = ttk.Frame(terminal_frame)
        toolbar.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Button(toolbar, text="清除屏幕", command=self.clear_screen).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(toolbar, text="重新连接", command=self.reconnect).pack(side=tk.LEFT, padx=(0, 5))
        
        # 终端显示区域
        terminal_container = ttk.Frame(terminal_frame, style='Terminal.TFrame')
        terminal_container.pack(fill=tk.BOTH, expand=True)
        
        self.terminal_text = scrolledtext.ScrolledText(
            terminal_container,
            bg='black',
            fg='white',
            insertbackground='white',
            selectbackground='#007acc',
            font=('Consolas', 10),
            wrap=tk.WORD,
            state=tk.DISABLED
        )
        self.terminal_text.pack(fill=tk.BOTH, expand=True, padx=1, pady=1)
        
        # 命令输入行
        input_frame = ttk.Frame(terminal_frame)
        input_frame.pack(fill=tk.X, pady=(5, 0))
        
        ttk.Label(input_frame, text=">").pack(side=tk.LEFT, padx=(0, 5))
        self.command_entry = ttk.Entry(input_frame)
        self.command_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 5))
        self.command_entry.bind('<Return>', self.send_command)
        
        ttk.Button(input_frame, text="发送", command=self.send_command).pack(side=tk.RIGHT)
        
    def create_status_bar(self, parent):
        """创建状态栏"""
        status_frame = ttk.Frame(parent, relief=tk.SUNKEN)
        status_frame.pack(fill=tk.X, pady=(5, 0))
        
        # 左侧状态信息
        self.connection_info = ttk.Label(status_frame, text="就绪")
        self.connection_info.pack(side=tk.LEFT, padx=5)
        
        # 右侧状态信息
        status_right = ttk.Frame(status_frame)
        status_right.pack(side=tk.RIGHT)
        
        self.cursor_pos = ttk.Label(status_right, text="行:1 列:1")
        self.cursor_pos.pack(side=tk.LEFT, padx=5)
        
        ttk.Label(status_right, text="|").pack(side=tk.LEFT, padx=2)
        
        self.encoding_label = ttk.Label(status_right, text="UTF-8")
        self.encoding_label.pack(side=tk.LEFT, padx=5)
        
        ttk.Label(status_right, text="|").pack(side=tk.LEFT, padx=2)
        
        self.time_label = ttk.Label(status_right, text=datetime.now().strftime("%H:%M:%S"))
        self.time_label.pack(side=tk.LEFT, padx=5)
        
        # 更新时间
        self.update_time()
        
    def strip_ansi_codes(self, text):
        """移除ANSI转义序列（颜色代码）"""
        # ANSI转义序列的正则表达式
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        return ansi_escape.sub('', text)
        
    def update_time(self):
        """更新时间显示"""
        self.time_label.config(text=datetime.now().strftime("%H:%M:%S"))
        self.root.after(1000, self.update_time)
        
    def toggle_connection(self):
        """切换连接状态"""
        if self.connected:
            self.disconnect()
        else:
            self.connect()
            
    def connect(self):
        """建立SSH连接"""
        host = self.host_entry.get().strip()
        port = self.port_entry.get().strip()
        username = self.username_entry.get().strip()
        password = self.password_entry.get()
        
        if not all([host, port, username]):
            messagebox.showerror("错误", "请填写完整的主机、端口和用户名信息")
            return
            
        if not port.isdigit():
            messagebox.showerror("错误", "端口号必须是数字")
            return
            
        self.connect_btn.config(state=tk.DISABLED)
        self.status_label.config(text="连接中...")
        self.connection_info.config(text=f"正在连接到 {host}:{port}")
        
        # 清空终端显示
        self.clear_screen()
        
        # 在新线程中连接
        threading.Thread(target=self._connect_thread, daemon=True).start()
        
    def _connect_thread(self):
        """连接线程"""
        try:
            host = self.host_entry.get().strip()
            port = int(self.port_entry.get().strip())
            username = self.username_entry.get().strip()
            password = self.password_entry.get()
            
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接参数
            connect_kwargs = {
                'hostname': host,
                'port': port,
                'username': username,
                'timeout': 15,
                'look_for_keys': False,
                'allow_agent': False
            }
            
            if self.auth_method.get() == "Password" and password:
                connect_kwargs['password'] = password
                
            self.ssh_client.connect(**connect_kwargs)
            
            # 连接成功
            self.connected = True
            self._update_ui_connected()
            self._append_output(f"*** 连接到 {host} 成功 ***\n\n")
            
            # 创建交互式shell
            self.shell = self.ssh_client.invoke_shell()
            self.shell.settimeout(0.1)
            self._start_receiving()
            
        except paramiko.ssh_exception.AuthenticationException:
            error_msg = "认证失败：用户名或密码错误"
            self._update_ui_error(error_msg)
            messagebox.showerror("认证失败", "请检查用户名和密码是否正确")
            
        except socket.timeout:
            error_msg = "连接超时：服务器未响应"
            self._update_ui_error(error_msg)
            messagebox.showerror("连接超时", "服务器未在指定时间内响应")
            
        except Exception as e:
            error_msg = f"连接失败：{str(e)}"
            self._update_ui_error(error_msg)
            messagebox.showerror("连接失败", error_msg)
            
    def _update_ui_connected(self):
        """更新UI为已连接状态"""
        self.root.after(0, lambda: self.status_label.config(text="已连接", style='Success.TLabel'))
        self.root.after(0, lambda: self.connect_btn.config(text="断开", state=tk.NORMAL))
        self.root.after(0, lambda: self.connection_info.config(text=f"已连接到 {self.host_entry.get()}"))
        
    def _update_ui_error(self, error_msg):
        """更新UI为错误状态"""
        self.root.after(0, lambda: self.status_label.config(text="连接失败", style='Error.TLabel'))
        self.root.after(0, lambda: self.connect_btn.config(state=tk.NORMAL))
        self.root.after(0, lambda: self.connection_info.config(text=error_msg))
        
    def _start_receiving(self):
        """开始接收数据"""
        threading.Thread(target=self._receive_thread, daemon=True).start()
        
    def _receive_thread(self):
        """接收数据线程"""
        while self.connected and self.shell:
            try:
                if self.shell.recv_ready():
                    data = self.shell.recv(4096).decode('utf-8', errors='ignore')
                    # 清理ANSI转义序列
                    clean_data = self.strip_ansi_codes(data)
                    self._append_output(clean_data)
                time.sleep(0.1)
            except socket.timeout:
                # 超时是正常的，继续循环
                continue
            except Exception as e:
                if self.connected:  # 如果仍然连接但出现异常，可能是连接断开
                    self._append_output(f"\n*** 连接异常: {str(e)} ***\n")
                    self.disconnect()
                break
                
    def _append_output(self, text):
        """添加输出到终端"""
        self.root.after(0, self._update_terminal, text)
        
    def _update_terminal(self, text):
        """更新终端显示"""
        self.terminal_text.config(state=tk.NORMAL)
        self.terminal_text.insert(tk.END, text)
        self.terminal_text.see(tk.END)
        self.terminal_text.config(state=tk.DISABLED)
        
    def send_command(self, event=None):
        """发送命令"""
        if not self.connected or not self.shell:
            messagebox.showwarning("警告", "未连接到SSH服务器")
            return
            
        command = self.command_entry.get().strip()
        if not command:
            return
            
        # 在终端显示命令（不显示ANSI代码）
        clean_command = self.strip_ansi_codes(command)
        self._append_output(f"\n$ {clean_command}\n")
        self.shell.send(command + '\n')
        self.command_entry.delete(0, tk.END)
        
    def disconnect(self):
        """断开连接"""
        self.connected = False
        if self.shell:
            try:
                self.shell.close()
            except:
                pass
        if self.ssh_client:
            try:
                self.ssh_client.close()
            except:
                pass
            
        self.status_label.config(text="未连接", style='Error.TLabel')
        self.connect_btn.config(text="连接", state=tk.NORMAL)
        self.connection_info.config(text="已断开连接")
        self._append_output("\n*** 连接已断开 ***\n")
        
    def clear_screen(self):
        """清空屏幕"""
        self.terminal_text.config(state=tk.NORMAL)
        self.terminal_text.delete(1.0, tk.END)
        self.terminal_text.config(state=tk.DISABLED)
        
    def reconnect(self):
        """重新连接"""
        if self.connected:
            self.disconnect()
        # 等待一下再连接
        self.root.after(500, self.connect)
        
    # 菜单功能实现
    def new_session(self):
        """新建会话"""
        # 清空所有字段
        self.host_entry.delete(0, tk.END)
        self.port_entry.delete(0, tk.END)
        self.port_entry.insert(0, "22")
        self.username_entry.delete(0, tk.END)
        self.password_entry.delete(0, tk.END)
        self.auth_method.set("Password")
        self.clear_screen()
        
    def open_session(self):
        """打开会话"""
        messagebox.showinfo("信息", "打开会话功能")
        
    def save_session(self):
        """保存会话"""
        host = self.host_entry.get().strip()
        if not host:
            messagebox.showerror("错误", "请输入主机地址")
            return
            
        session_name = f"{host}_{self.username_entry.get()}"
        self.sessions[session_name] = {
            'host': host,
            'port': self.port_entry.get(),
            'username': self.username_entry.get(),
            'auth_method': self.auth_method.get()
        }
        messagebox.showinfo("成功", f"会话 '{session_name}' 已保存")
        
    def on_session_change(self, event):
        """会话切换"""
        # 这个功能在简化版中暂时不实现
        pass
        
    def copy_text(self):
        """复制文本"""
        try:
            selected_text = self.terminal_text.get(tk.SEL_FIRST, tk.SEL_LAST)
            self.root.clipboard_clear()
            self.root.clipboard_append(selected_text)
        except:
            pass
            
    def paste_text(self):
        """粘贴文本"""
        try:
            clipboard_text = self.root.clipboard_get()
            self.command_entry.insert(tk.INSERT, clipboard_text)
        except:
            pass
            
    def select_all(self):
        """全选"""
        self.terminal_text.focus_set()
        self.terminal_text.tag_add(tk.SEL, "1.0", tk.END)
        
    def toggle_theme(self):
        """切换主题"""
        messagebox.showinfo("主题", "主题切换功能")
        
    def font_settings(self):
        """字体设置"""
        messagebox.showinfo("字体", "字体设置功能")
        
    def session_manager(self):
        """会话管理器"""
        messagebox.showinfo("会话管理器", "会话管理功能")
        
    def options(self):
        """选项设置"""
        messagebox.showinfo("选项", "选项设置功能")

def main():
    try:
        import paramiko
    except ImportError:
        print("错误: 未找到paramiko库，请先安装:")
        print("pip install paramiko")
        return
        
    root = tk.Tk()
    app = XShellLikeSSHClient(root)
    root.mainloop()

if __name__ == "__main__":
    main()